#include "vkcompute.h"
#include "vkutils.h"

namespace vkinfer
{
    VkCompute::VkCompute(VkDevice device, VkQueue queue, uint32_t queue_index)
    {
        this->device = device;
        this->queue = queue;

        VkCommandPoolCreateInfo commandPoolCreateInfo = {};
        commandPoolCreateInfo.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
        commandPoolCreateInfo.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT;
        commandPoolCreateInfo.queueFamilyIndex = queue_index;

        VK_CHECK_RESULT(vkCreateCommandPool(device, &commandPoolCreateInfo, NULL, &commandPool));

        VkCommandBufferAllocateInfo commandBufferAllocateInfo = {};
        commandBufferAllocateInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
        commandBufferAllocateInfo.commandPool = commandPool;
        commandBufferAllocateInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
        commandBufferAllocateInfo.commandBufferCount = 1;
        VK_CHECK_RESULT(vkAllocateCommandBuffers(device, &commandBufferAllocateInfo, &commandBuffer));
    }

    VkCompute::~VkCompute()
    {
        vkDestroyCommandPool(device, commandPool, nullptr);
    }

    void VkCompute::begin_command_buffer()
    {
        vkDeviceWaitIdle(device);
        VkCommandBufferBeginInfo beginInfo = {};
        beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
        beginInfo.flags = VK_COMMAND_BUFFER_USAGE_SIMULTANEOUS_USE_BIT;
        VK_CHECK_RESULT(vkBeginCommandBuffer(commandBuffer, &beginInfo)); // start recording commands.
    }

    void VkCompute::end_command_buffer()
    {
        VK_CHECK_RESULT(vkEndCommandBuffer(commandBuffer));
    }

    void VkCompute::submit(VkFence fence)
    {
        VkSubmitInfo submitInfo = {};
        submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
        submitInfo.commandBufferCount = 1;
        submitInfo.pCommandBuffers = &commandBuffer;

        VK_CHECK_RESULT(vkQueueSubmit(queue, 1, &submitInfo, fence));
    }

    void VkCompute::buffer_barrier(const VkDependencyInfo& dep)
    {
        vkCmdPipelineBarrier2(commandBuffer, &dep);
    }

    void VkCompute::barrier_writeread(const std::vector<std::shared_ptr<VkBlob>>& buffers, bool top_layer)
    {
#if 1
        if (top_layer) return;
        for (size_t i = 0; i < buffers.size(); ++i)
        {
            VkBufferMemoryBarrier* barriers = new VkBufferMemoryBarrier[1];
            barriers[0].sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER;
            barriers[0].pNext = 0;
            barriers[0].srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT;
            barriers[0].dstAccessMask = VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT;
            barriers[0].srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
            barriers[0].dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
            barriers[0].buffer = buffers[i]->buffer;
            barriers[0].offset = 0;
            barriers[0].size = buffers[i]->size;

            VkPipelineStageFlags src_stage = VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT;
            VkPipelineStageFlags dst_stage = VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT;
            vkCmdPipelineBarrier(commandBuffer, src_stage, dst_stage, 0, 0, 0, 1, barriers, 0, 0);
            delete[] barriers;
        }
#endif
    }
}
